from PIL import Image, ImageEnhance, ImageFilter
import os


def augment_image(image_path, output_folder):
    # 确保 output_folder 存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 打开图片
    image = Image.open(image_path)

    # 提取不带路径的文件名和扩展名
    base_filename = os.path.basename(image_path)
    filename, file_extension = os.path.splitext(base_filename)

    # 检查是否已增强
    enhanced_suffixes = ['_flipped', '_rotated', '_scaled', '_contrast', '_noisy']
    if any(filename.endswith(suffix) for suffix in enhanced_suffixes):
        print(f"Skipping already augmented image: {filename}")
        return

    # 翻转
    flipped_image = image.transpose(Image.FLIP_LEFT_RIGHT)
    flipped_image.save(os.path.join(output_folder, f'{filename}_flipped{file_extension}'))

    # 旋转
    rotated_image = image.rotate(45, expand=True)
    rotated_image.save(os.path.join(output_folder, f'{filename}_rotated{file_extension}'))

    # 缩放
    scaled_image = image.resize((int(image.width * 0.5), int(image.height * 0.5)))
    scaled_image.save(os.path.join(output_folder, f'{filename}_scaled{file_extension}'))

    # 调整对比度
    enhancer = ImageEnhance.Contrast(image)
    contrast_image = enhancer.enhance(2.0)
    contrast_image.save(os.path.join(output_folder, f'{filename}_contrast{file_extension}'))

    # 添加高斯噪声
    noisy_image = image.filter(ImageFilter.GaussianBlur(radius=1))
    noisy_image.save(os.path.join(output_folder, f'{filename}_noisy{file_extension}'))


# 图片存放的文件夹路径
image_folder = 'imgs'

# 存放增强后图片的文件夹路径
output_folder = 'imgs'

# 对 image_folder 中的图片进行增强
for filename in os.listdir(image_folder):
    if filename.endswith(('.jpg', '.jpeg', '.png')):
        image_path = os.path.join(image_folder, filename)
        augment_image(image_path, output_folder)